# 对一副图像加噪，进行几何均值，算术均值，谐波，逆谐波处理，显示图像的结果；


import cv2
import numpy as np
import matplotlib.pyplot as plt
import scipy
import scipy.stats
import random

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def GaussieNoisy(image, sigma):
    # 高斯噪声
    img = image.astype(np.int16)  # 此步是为了避免像素点小于0，大于255的情况
    mu = 0
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            for k in range(img.shape[2]):
                img[i, j, k] = img[i, j, k] + random.gauss(mu=mu, sigma=sigma)
    img[img > 255] = 255
    img[img < 0] = 0
    img = img.astype(np.uint8)
    return img


def spNoisy(image, s_vs_p=0.5, amount=0.004):
    # 椒盐噪声
    out = np.copy(image)
    num_salt = np.ceil(amount * image.size * s_vs_p)
    coords = [np.random.randint(0, i - 1, int(num_salt)) for i in image.shape]
    out[tuple(coords)] = 1
    num_pepper = np.ceil(amount * image.size * (1. - s_vs_p))
    coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in image.shape]
    out[tuple(coords)] = 0
    return out


def ArithmeticMeanAlogrithm(image):
    # 算术均值滤波
    new_image = np.zeros(image.shape)
    image = cv2.copyMakeBorder(image, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
    for i in range(1, image.shape[0] - 1):
        for j in range(1, image.shape[1] - 1):
            new_image[i - 1, j - 1] = np.mean(image[i - 1:i + 2, j - 1:j + 2])
    new_image = (new_image - np.min(image)) * (255 / np.max(image))
    return new_image.astype(np.uint8)


def rgbArithmeticMean(image):
    # 三个通道都进行算术均值滤波，最后合并
    r, g, b = cv2.split(image)
    r = ArithmeticMeanAlogrithm(r)
    g = ArithmeticMeanAlogrithm(g)
    b = ArithmeticMeanAlogrithm(b)
    return cv2.merge([r, g, b])


def GeometricMeanOperator(roi):
    roi = roi.astype(np.float64)
    p = np.prod(roi)
    return p ** (1 / (roi.shape[0] * roi.shape[1]))


def GeometricMeanAlogrithm(image):
    # 几何均值滤波
    new_image = np.zeros(image.shape)
    image = cv2.copyMakeBorder(image, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
    for i in range(1, image.shape[0] - 1):
        for j in range(1, image.shape[1] - 1):
            new_image[i - 1, j - 1] = GeometricMeanOperator(image[i - 1:i + 2, j - 1:j + 2])
    new_image = (new_image - np.min(image)) * (255 / np.max(image))
    return new_image.astype(np.uint8)


def rgbGemotriccMean(image):
    # 三个通道都进行几何均值滤波，最后合并
    r, g, b = cv2.split(image)
    r = GeometricMeanAlogrithm(r)
    g = GeometricMeanAlogrithm(g)
    b = GeometricMeanAlogrithm(b)
    return cv2.merge([r, g, b])


def HarmonicMeanOperator(roi):
    roi = roi.astype(np.float64)
    if 0 in roi:
        roi = 0
    else:
        roi = scipy.stats.hmean(roi.reshape(-1))
    return roi


def HarmonicMeanAlogrithm(image):
    # 谐波均值滤波
    new_image = np.zeros(image.shape)
    image = cv2.copyMakeBorder(image, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
    for i in range(1, image.shape[0] - 1):
        for j in range(1, image.shape[1] - 1):
            new_image[i - 1, j - 1] = HarmonicMeanOperator(image[i - 1:i + 2, j - 1:j + 2])
    new_image = (new_image - np.min(image)) * (255 / np.max(image))
    return new_image.astype(np.uint8)


def rgbHarmonicMean(image):
    r, g, b = cv2.split(image)
    r = HarmonicMeanAlogrithm(r)
    g = HarmonicMeanAlogrithm(g)
    b = HarmonicMeanAlogrithm(b)
    return cv2.merge([r, g, b])


def Contra_harmonicMeanOperator(roi, q):
    roi = roi.astype(np.float64)
    return np.mean((roi) ** (q + 1)) / np.mean((roi) ** (q))


def Contra_harmonicMeanAlogrithm(image, q):
    # 逆谐波均值滤波
    new_image = np.zeros(image.shape)
    image = cv2.copyMakeBorder(image, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
    for i in range(1, image.shape[0] - 1):
        for j in range(1, image.shape[1] - 1):
            new_image[i - 1, j - 1] = Contra_harmonicMeanOperator(image[i - 1:i + 2, j - 1:j + 2], q)
    new_image = (new_image - np.min(image)) * (255 / np.max(image))
    return new_image.astype(np.uint8)


def rgbContra_harmonicMean(image, q):
    r, g, b = cv2.split(image)
    r = Contra_harmonicMeanAlogrithm(r, q)
    g = Contra_harmonicMeanAlogrithm(g, q)
    b = Contra_harmonicMeanAlogrithm(b, q)
    return cv2.merge([r, g, b])


if __name__ == '__main__':
    img = cv2.imread("img/lena.png")
    img = cv2.resize(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), (200, 200))
    img_gaussie_noisy = GaussieNoisy(img, 18)
    img_sp_noisy = spNoisy(img)

    plt.subplot(161), plt.imshow(img), plt.title("原图"), plt.xticks([]), plt.yticks([])
    plt.subplot(162), plt.imshow(img_gaussie_noisy), plt.title("加高斯噪声"), plt.xticks(
        []), plt.yticks([])
    plt.subplot(163), plt.imshow(rgbArithmeticMean(img_gaussie_noisy)), plt.title("算数均值"), plt.xticks(
        []), plt.yticks([])
    plt.subplot(164), plt.imshow(rgbGemotriccMean(img_gaussie_noisy)), plt.title("几何均值"), plt.xticks(
        []), plt.yticks([])
    plt.subplot(165), plt.imshow(rgbHarmonicMean(img_gaussie_noisy)), plt.title("谐波"), plt.xticks(
        []), plt.yticks([])
    plt.subplot(166), plt.imshow(rgbContra_harmonicMean(img_gaussie_noisy, 2)), plt.title("逆谐波"), plt.xticks(
        []), plt.yticks([])
    plt.show()
    plt.subplot(161), plt.imshow(img), plt.title("原图"), plt.xticks([]), plt.yticks([])
    plt.subplot(162), plt.imshow(img_sp_noisy), plt.title("加椒盐噪声"), plt.xticks(
        []), plt.yticks([])
    plt.subplot(163), plt.imshow(rgbArithmeticMean(img_sp_noisy)), plt.title("算数均值"), plt.xticks(
        []), plt.yticks([])
    plt.subplot(164), plt.imshow(rgbGemotriccMean(img_sp_noisy)), plt.title("几何均值"), plt.xticks(
        []), plt.yticks([])
    plt.subplot(165), plt.imshow(rgbHarmonicMean(img_sp_noisy)), plt.title("谐波"), plt.xticks(
        []), plt.yticks([])
    plt.subplot(166), plt.imshow(rgbContra_harmonicMean(img_sp_noisy, 2)), plt.title("逆谐波"), plt.xticks(
        []), plt.yticks([])
    plt.show()

